@echo off
REM PRISM 高性能训练 - 快速启动脚本 (Windows)

echo ==================================================
echo     PRISM 高性能训练 - 快速启动
echo ==================================================
echo.

REM 检查Python环境
echo [1/5] 检查Python环境...
python --version >nul 2>&1
if errorlevel 1 (
    echo 错误: 未找到Python
    pause
    exit /b 1
)
echo √ Python已安装
echo.

REM 检查依赖包
echo [2/5] 检查依赖包...
python -c "import torch, ultralytics, optuna" >nul 2>&1
if errorlevel 1 (
    echo 警告: 部分依赖包未安装
    echo 运行: pip install -r requirements.txt
) else (
    echo √ 依赖包已安装
)
echo.

REM 检查数据集
echo [3/5] 检查数据集...
if exist "dataset" (
    echo √ 数据集目录存在
) else (
    echo 错误: 未找到dataset目录
    pause
    exit /b 1
)
echo.

REM 检查GPU
echo [4/5] 检查GPU...
python -c "import torch; print(f'GPU可用: {torch.cuda.is_available()}'); print(f'GPU数量: {torch.cuda.device_count()}')"
echo.

REM 创建必要的目录
echo [5/5] 创建输出目录...
if not exist "high_performance_results" mkdir high_performance_results
if not exist "optimization_results" mkdir optimization_results
if not exist "weights" mkdir weights
echo √ 目录创建完成
echo.

echo ==================================================
echo     环境检查完成，选择训练模式:
echo ==================================================
echo.
echo   1) 完整训练流程（推荐，12-24小时）
echo   2) 仅Stage 1训练（2-4小时）
echo   3) 仅Stage 2训练（3-6小时）
echo   4) 仅超参数优化（6-12小时）
echo   5) 仅评估（10分钟）
echo   6) 快速测试（所有阶段，少量epoch）
echo   0) 退出
echo.
set /p choice="请选择 [0-6]: "

if "%choice%"=="1" (
    echo 启动完整训练流程...
    python train_high_performance.py --mode full
) else if "%choice%"=="2" (
    echo 启动Stage 1训练...
    python train_high_performance.py --mode stage1
) else if "%choice%"=="3" (
    echo 启动Stage 2训练...
    python train_high_performance.py --mode stage2
) else if "%choice%"=="4" (
    echo 启动超参数优化...
    set /p n_trials="输入试验次数（建议30-50）: "
    if "%n_trials%"=="" set n_trials=30
    python train_high_performance.py --mode optimize --n-trials %n_trials%
) else if "%choice%"=="5" (
    echo 启动评估...
    python train_high_performance.py --mode evaluate
) else if "%choice%"=="6" (
    echo 快速测试模式（减少epoch）...
    REM 创建临时配置
    (
        echo stage1:
        echo   epochs: 10
        echo   batch_size: 16
        echo   img_size: 640
        echo   patience: 5
        echo.
        echo stage2:
        echo   epochs: 10
        echo   batch_size: 32
        echo   learning_rate: 0.0001
        echo   patience: 5
        echo   model:
        echo     use_fpn: false
        echo     use_roi_align: false
        echo.
        echo optimization:
        echo   n_trials: 5
        echo   timeout: null
    ) > config_quick_test.yaml
    python train_high_performance.py --mode full --config config_quick_test.yaml
) else if "%choice%"=="0" (
    echo 退出
    exit /b 0
) else (
    echo 无效选择
    pause
    exit /b 1
)

echo.
echo ==================================================
echo     训练完成！
echo ==================================================
echo.
echo 查看结果:
echo   - 训练历史: high_performance_results\training_history.json
echo   - 最终评估: high_performance_results\final_evaluation\
echo   - 优化结果: optimization_results\
echo   - 日志文件: training_*.log
echo.
pause
